module im2col #(
    parameter int unsigned P_CH    = 4,
    parameter int unsigned N_CH    = 16,
    parameter int unsigned N_IH    = 8,
    parameter int unsigned N_IW    = 8,
    parameter int unsigned K       = 3,
    parameter int unsigned P       = 1,
    parameter int unsigned S       = 1,
    parameter int unsigned A_BIT   = 8,
    parameter int unsigned PAD_VAL = 0
) (
    input  logic                  clk,
    input  logic                  rst_n,
    input  logic [P_CH*A_BIT-1:0] in_data,
    input  logic                  in_valid,
    output logic                  in_ready,
    output logic [P_CH*A_BIT-1:0] out_data,
    output logic                  out_valid,
    input  logic                  out_ready
);

    localparam int unsigned N_OH = (N_IH + 2 * P - K) / S + 1;
    localparam int unsigned N_OW = (N_IW + 2 * P - K) / S + 1;
    localparam int unsigned FOLD = N_CH / P_CH;
    localparam int unsigned PRELOAD_PIXELS = K - 1;

    typedef enum logic [1:0] {
        ST_INIT,
        ST_PROC
    } state_t;
    state_t                           state;
    logic        [    P_CH*A_BIT-1:0] line_buffer       [K][N_IW][FOLD];
    logic        [  $clog2(N_IW+K):0] cntr_init_pixels;
    logic        [  $clog2(FOLD+1):0] cntr_init_fold;
    logic        [  $clog2(N_IH+1):0] cntr_ih;
    logic        [  $clog2(N_IW+1):0] cntr_iw;
    logic        [  $clog2(FOLD+1):0] cntr_fold;
    logic        [     $clog2(K+1):0] cntr_kh;
    logic        [     $clog2(K+1):0] cntr_kw;
    logic                             pipe_en_in;
    logic                             pipe_en_out;
    logic                             pipe_en;
    logic                             need_read_input;
    logic        [  $clog2(N_IH+1):0] ih_idx_to_write;
    logic signed [$clog2(N_IH+K)+1:0] ih_idx_to_read;
    logic        [  $clog2(N_IW+K):0] iw_idx_to_write;
    logic        [  $clog2(N_IW+K):0] iw_idx_to_write_raw;
    logic signed [$clog2(N_IW+K)+1:0] iw_idx_to_read;
    logic        [       $clog2(K):0] line_idx_to_write;
    logic        [       $clog2(K):0] line_idx_to_read;
    logic        [  $clog2(N_OH+1):0] oh_idx;
    logic        [  $clog2(N_OW+1):0] ow_idx;
    logic                             is_padding;
    logic                             need_write_output;
    logic        [    P_CH*A_BIT-1:0] pad_data;

    always_comb begin
        for (int i = 0; i < P_CH; i++) begin
            pad_data[i*A_BIT+:A_BIT] = PAD_VAL;
        end
    end

    assign pipe_en_in = need_read_input ? in_valid : 1'b1;
    assign pipe_en_out = out_ready;
    assign pipe_en = pipe_en_in && pipe_en_out;
    assign in_ready = (state == ST_INIT) || (need_read_input && out_ready);
    assign out_valid = pipe_en_in && need_write_output;
    assign out_data = is_padding ? pad_data : line_buffer[line_idx_to_read][iw_idx_to_read][cntr_fold];
    assign need_read_input = (state == ST_PROC) && (ih_idx_to_write < N_IH) && (cntr_kh == 0) && (cntr_kw == 0);
    assign need_write_output = (state == ST_PROC) && (cntr_ih % S == 0) && (cntr_iw % S == 0);
    assign oh_idx = cntr_ih / S;
    assign ow_idx = cntr_iw / S;
    assign ih_idx_to_read = oh_idx * S + cntr_kh - P;
    assign iw_idx_to_read = ow_idx * S + cntr_kw - P;
    assign is_padding        = (ih_idx_to_read < 0) || (ih_idx_to_read >= N_IH) || 
                               (iw_idx_to_read < 0) || (iw_idx_to_read >= N_IW);
    assign line_idx_to_read = ih_idx_to_read % K;
    assign line_idx_to_write = ih_idx_to_write % K;
    assign iw_idx_to_write_raw = cntr_iw + PRELOAD_PIXELS;
    assign iw_idx_to_write = (iw_idx_to_write_raw >= N_IW) ? (iw_idx_to_write_raw - N_IW) : iw_idx_to_write_raw;

    always_ff @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            state            <= ST_INIT;
            cntr_init_pixels <= 0;
            cntr_init_fold   <= 0;
            cntr_ih          <= 0;
            cntr_iw          <= 0;
            cntr_fold        <= 0;
            cntr_kh          <= 0;
            cntr_kw          <= 0;
            ih_idx_to_write  <= 0;
            for (int h = 0; h < K; h++) begin
                for (int w = 0; w < N_IW; w++) begin
                    for (int f = 0; f < FOLD; f++) begin
                        line_buffer[h][w][f] <= '0;
                    end
                end
            end
        end else begin
            case (state)
                ST_INIT: begin
                    if (in_valid) begin
                        if (cntr_init_pixels < N_IW) begin
                            line_buffer[0][cntr_init_pixels][cntr_init_fold] <= in_data;
                        end else begin
                            line_buffer[1][cntr_init_pixels-N_IW][cntr_init_fold] <= in_data;
                        end
                        if (cntr_init_fold == FOLD - 1) begin
                            cntr_init_fold <= 0;
                            if (cntr_init_pixels == N_IW + PRELOAD_PIXELS - 1) begin
                                state           <= ST_PROC;
                                ih_idx_to_write <= 1;
                            end else begin
                                cntr_init_pixels <= cntr_init_pixels + 1;
                            end
                        end else begin
                            cntr_init_fold <= cntr_init_fold + 1;
                        end
                    end
                end

                ST_PROC: begin
                    if (in_valid && in_ready) begin
                        line_buffer[line_idx_to_write][iw_idx_to_write][cntr_fold] <= in_data;
                    end
                    if (pipe_en) begin
                        if (cntr_kw == K - 1) begin
                            cntr_kw <= 0;
                            if (cntr_kh == K - 1) begin
                                cntr_kh <= 0;
                                if (cntr_fold == FOLD - 1) begin
                                    cntr_fold <= 0;
                                    if (iw_idx_to_write_raw == N_IW - 1) begin
                                        ih_idx_to_write <= ih_idx_to_write + 1;
                                    end
                                    if (cntr_iw == N_IW - 1) begin
                                        cntr_iw <= 0;
                                        if (cntr_ih == N_IH - 1) begin
                                            state            <= ST_INIT;
                                            cntr_ih          <= 0;
                                            cntr_init_pixels <= 0;
                                            cntr_init_fold   <= 0;
                                            ih_idx_to_write  <= 1;
                                        end else begin
                                            cntr_ih <= cntr_ih + 1;
                                        end
                                    end else begin
                                        cntr_iw <= cntr_iw + 1;
                                    end
                                end else begin
                                    cntr_fold <= cntr_fold + 1;
                                end
                            end else begin
                                cntr_kh <= cntr_kh + 1;
                            end
                        end else begin
                            cntr_kw <= cntr_kw + 1;
                        end
                    end
                end
            endcase
        end
    end
endmodule
